import monai.transforms
import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np

from skimage.measure import shannon_entropy
from einops import rearrange

import SimpleITK as sitk
from time import time


class GridMaskLayer(nn.Module):
    """
    对输入图像进行动态遮挡
    """
    def __init__(self, dynamic_masking=True, hierarchical_masking=None, basic_mask_ratio=0.5,
                 scale=0.5, drop_ratio=0.3):
        super(GridMaskLayer, self).__init__()
        self.pre_defined_grid_resolution = [32]

        # mask ratio 缩放尺度
        self.scale = scale
        self.drop_ratio = drop_ratio
        self.step = 0

        self.dynamic_masking = dynamic_masking
        self.hierarchical_masking = None if hierarchical_masking == 0. else hierarchical_masking
        self.basic_mask_ratio = basic_mask_ratio

    def forward(self, x):
        x = self.mask_token(x)
        return x

    def initial(self, x):
        self.grid_resolution = self.pre_defined_grid_resolution
        self.mask_list = []

    def mask_token(self, x):
        # token resolution
        b_, _, _, _, w_ = x.size()

        for batch_index in range(b_):
            mask_list = []
            if self.grid_resolution[0] > w_:
                pad = self.grid_resolution[0] % w_ // 2
            else:
                pad = w_ % self.grid_resolution[0] // 2

            if pad != 0:
                padding = nn.ConstantPad3d(int(pad), 0.)
                img = padding(x[batch_index].unsqueeze(0)).squeeze(0)
            else:
                img = x[batch_index]

            n, d, h, w = img.size()

            gr = self.grid_resolution[0]

            img = img.view(n, gr, d // gr, gr, h // gr, gr, w // gr)
            img = img.permute(0, 2, 4, 6, 1, 3, 5)

            # [n, grid_num, token_num_in_grid]
            img = rearrange(img, 'n d h w res_d res_h res_w -> n (d h w) (res_d res_h res_w)')

            num_mask = int(img.size()[-1] * self.basic_mask_ratio)

            mask_seq = np.concatenate([np.zeros(num_mask),
                                       np.ones(img.size()[-1] - num_mask)])

            np.random.shuffle(mask_seq)
            mask_list = torch.from_numpy(np.array(mask_seq))
            mask_seq = torch.from_numpy(np.array(mask_seq))
            
            img *= mask_seq.to(img.device)

            img = img.reshape(n, -1, gr, gr, gr)
            img = img.reshape(n, d // gr, h // gr, w // gr, gr, gr, gr)
            img = img.permute(0, 4, 1, 5, 2, 6, 3)
            img = rearrange(img, 'n gz nz gh nh gw nw -> n (gz nz) (gh nh) (gw nw)')

            if pad != 0:
                img = img[:, pad:-pad, pad:-pad, pad:-pad]

            x[batch_index] = img
            self.mask_list.append(mask_list)
        return x

    def reset_step(self):
        self.step = 0
        # [b]
        self.grid_resolution = []
        # [b, grid_num]
        self.shannon_entropy = []
        # [b, grid_num]
        self.dynamic_mask_ratio = []

    def get_masked_gt(self, x):
        b, c, d, h, w = x.size()

        # patchy
        x = x.view(b, c, d // 4, 4, h // 4, 4, w // 4, 4)
        x = x.permute(0, 1, 2, 4, 6, 3, 5, 7)
        x = rearrange(x, 'b c res_z res_h res_w z h w -> b c (res_z res_h res_w) (z h w)')

        for batch_index in range(b):
            image = x[batch_index]

            image = image.permute(0, 2, 1)
            
            image *= (1. - self.mask_list[0]).to(image.device)

            image = image.permute(0, 2, 1)
            
            x[batch_index] = image

        x = x.view(b, c, d // 4, h // 4, w // 4, 4, 4, 4).permute(0, 1, 2, 5, 3, 6, 4, 7)
        x = rearrange(x, 'b c res_z z res_h h res_w w -> b c (res_z z) (res_h h) (res_w w)')
        return x






